import os
import pandas as pd
import re
from openpyxl import load_workbook
from openpyxl.styles import PatternFill


class DataCleaner:
    """
    数据清洗类，用于批量处理Excel文件，筛选无效数据并标红被删除的行。
    """

    def __init__(self, raw_data_folder, clean_data_folder):
        """
        初始化DataCleaner对象。
        :param raw_data_folder: 原始数据文件夹路径。
        :param clean_data_folder: 清洗后数据文件夹路径。
        """
        self.raw_data_folder = raw_data_folder
        self.clean_data_folder = clean_data_folder
        os.makedirs(clean_data_folder, exist_ok=True)

    @staticmethod
    def is_invalid_text(text):
        """
        判断文本是否为无效内容。
        :param text: 文本内容字符串。
        :return: 如果是无效内容返回True，否则返回False。
        """
        useless_patterns = [
            r"^嗯+.*",        # 语气词开头：例如“嗯”“嗯嗯”
            r"^啊+.*",        # 语气词开头：例如“啊”“啊啊”
            r"^哦+.*",        # 语气词开头：例如“哦”“哦哦”
            r"^哈+.*",        # 语气词开头：例如“哈哈”“哈哈哈”
            r"^\s*$"          # 空白行或空字符串
        ]
        for pattern in useless_patterns:
            if re.match(pattern, text):
                return True
        return False

    def process_file(self, file_path):
        """
        处理单个Excel文件，筛选无效数据并保存结果。
        :param file_path: 原始文件的路径。
        """
        # 读取Excel文件
        df = pd.read_excel(file_path, header=None, names=["时间", "文本内容", "场景", "标签"])
        df["status"] = "保留"  # 默认保留所有数据

        # 筛选无效数据
        for index, row in df.iterrows():
            text = row["文本内容"]
            if self.is_invalid_text(text):  # 如果文本是无效内容
                df.at[index, "status"] = "删除"

        # 修改列头为最终格式
        cleaned_df = df.rename(columns={"文本内容": "text", "标签": "label"})

        # 保存清洗后的数据
        filename = os.path.basename(file_path)
        output_path = os.path.join(self.clean_data_folder, filename)
        cleaned_df.to_excel(output_path, index=False)

        # 标红被删除的行
        self.mark_deleted_rows(output_path, cleaned_df)

    @staticmethod
    def mark_deleted_rows(file_path, cleaned_df):
        """
        标红被删除的行。
        :param file_path: 清洗后文件的路径。
        :param cleaned_df: 清洗后的数据框。
        """
        # 加载Excel文件
        workbook = load_workbook(file_path)
        sheet = workbook.active

        # 设置标红格式
        red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")
        for row_idx, status in enumerate(cleaned_df["status"], start=2):  # 跳过表头
            if status == "删除":
                for col_idx in range(1, sheet.max_column + 1):
                    sheet.cell(row=row_idx, column=col_idx).fill = red_fill

        # 保存文件
        workbook.save(file_path)

    def process_all_files(self):
        """
        批量处理raw_data_folder中的所有Excel文件。
        """
        for filename in os.listdir(self.raw_data_folder):
            if filename.endswith('.xlsx'):  # 确保处理的是Excel文件
                file_path = os.path.join(self.raw_data_folder, filename)
                print(f"正在处理文件: {filename}")
                self.process_file(file_path)
        print("所有数据清洗完成，结果已保存到 clean_data 文件夹中。")


# 主程序入口
if __name__ == "__main__":
    # 定义文件夹路径
    RAW_DATA_FOLDER = 'raw_data/'  # 原始数据文件夹
    CLEAN_DATA_FOLDER = 'clean_data/'  # 清洗后数据文件夹

    # 创建DataCleaner对象并处理所有文件
    cleaner = DataCleaner(RAW_DATA_FOLDER, CLEAN_DATA_FOLDER)
    cleaner.process_all_files()
